#include <iostream>
#include <fstream>
#include <vector>
#include <string>
#include <complex>
#include <math.h>
#include <algorithm>
#include <opencv2/opencv.hpp>
#include "wavelet.h"
//////////////////////////////////////////////////////////////////////////
using namespace cv;
//Daubechies 2
double Ld[] = {-0.1294,   0.2241,   0.8365,   0.4830};
double Hd[] = {-0.4830,   0.8365,  -0.2241,  -0.1294};
double Lr[] = { 0.4830,   0.8365,   0.2241,  -0.1294};
double Hr[] = {-0.1294,  -0.2241,   0.8365,  -0.4830};
//////////////////////////////////////////////////////////////////////////
/*
 signal extension for decomposition. the extension operated as following:
 ABCDEFG ABC...
*/
int signal_extension_d( double * original, int n, double * extended, int m, int type )
{
	int i;

	if ( n < m ) return( 0 );  // n should lager than or equal m, otherwise no extension

	switch( type )
	{
	case PERIODIC_EXTENSION:
	default:  
		for( i = 0; i < n; i++ )
			extended[i] = original[i];
		for ( i = 0; i < m-1; i++ )
			extended[n+i] = original[i];
		break;	
	}

	return( n+m-1 );
}
//////////////////////////////////////////////////////////////////////////
/*
 signal extension for reconstruction. the extension operated as following:
 ... EFG ABCDEFG
*/
int signal_extension_r( double * original, int n, double * extended, int m, int type )
{
	int i;

	if ( n < m ) return( 0 );  // n should lager than or equal m, otherwise no extension

	switch( type )
	{
	case PERIODIC_EXTENSION:
	default:  // default is PERIODIC_EXTENSION
		for( i = 0; i < m-1; i++ )
			extended[i] = original[n-m+i+1];
		for ( i = 0; i < n; i++ )
			extended[m-1+i] = original[i];
		break;
	}

	return( n+m-1 );
}
//////////////////////////////////////////////////////////////////////////
/*
  1D convolution.
  parameters:
  double * data     --- original 1D signal
  int n             --- dimension of the original signal
  double * filter   --- coefficients of filter
  int m             --- length of convolution filter
  int type          --- extension type, three types of extension are supported
  double * result   --- the convoluted result, should have the same length with original data
  return value      --- length of result signal
*/
int convolution( double * data, int n, double * filter, int m, int type, double * result, int op )
{
	int mn;
	double * data1 = NULL;
	double sum;
	int i, j;

	mn = m+n-1;
	data1 = new double[mn];
	if ( data1 == NULL ) return( 0 );

	if ( op == DECOMPOSITION )
		mn = signal_extension_d( data, n, data1, m, type );
	else
		mn = signal_extension_r( data, n, data1, m, type );
	if ( mn > 0 )
	{
		for ( i = 0; i < n; i++ )
		{
			sum = 0;
			for ( j = 0; j < m; j++ )
				sum = data1[i+j]*filter[m-1-j]+sum;
			result[i] = sum;
		}
		mn = n;
	}
	delete data1;

	return( mn );
}
//////////////////////////////////////////////////////////////////////////
/*
 down-sample data, i.e., just keep one sample for per two data in a sequence.
 double * data        --- original data
 int n                --- data points of original data
 double * downdata    --- down-sampled data, should be of half the lenght of the original sequence
*/
int downsampling( double * data, int n, double * downdata )
{
	int i, m;

	m = n/2;
	for( i = 0; i < m; i++ )
		downdata[i] = data[i*2];

	return( m );
}

/*
 interpolate data, i.e., insert one zero between each two samples in a sequence.
 double * data        --- original data
 int n                --- data points of original data
 double * intpdata    --- interpolated data, should be of double the lenght of the original sequence
*/
int interpolation( double * data, int n, double * intpdata )
{
	int i;

	for( i = 0; i < n; i++ )
	{
		intpdata[2*i] = data[i];
		intpdata[2*i+1] = 0;
	}

	return( 2*n );
}
//////////////////////////////////////////////////////////////////////////
/*
 one round wavelet decomposition
*/
int wavelet_analysis( double * data, int n,   // original data
					  double * Hi, int n_Hi,  // high pass filter
				      double * Lo, int n_Lo,  // low pass filter
				      double * approximation, // approximation coefficients
				      double * detail,        // detail coefficients
				      int type,               // extension type
				      int ifsubsample )       // determine if do sub-sampling
{
	int m;
	double * approximation1, * detail1;

	if ( ifsubsample == SUBSAMPLING )
	{
		approximation1 = new double[n];
		detail1 = new double[n];
		m = convolution( data, n, Hi, n_Hi, type, detail1, DECOMPOSITION );
		m = convolution( data, n, Lo, n_Lo, type, approximation1, DECOMPOSITION );
		m = downsampling( detail1, n, detail );
		m = downsampling( approximation1, n, approximation );
		delete approximation1;
		delete detail1;
	}
	else
	{
		m = convolution( data, n, Hi, n_Hi, type, detail, DECOMPOSITION );
		m = convolution( data, n, Lo, n_Lo, type, approximation, DECOMPOSITION );
	}

	return( m );
}
//////////////////////////////////////////////////////////////////////////
/*
 one round wavelet synthesis
*/
int wavelet_synthesis( double * approximation, // approximation coefficients
				       double * detail, int n, // detail coefficients 
					   double * Hi, int n_Hi,  // high pass filter
				       double * Lo, int n_Lo,  // low pass filter
				       double * data,          // synthesis data
				       int type,               // extension type
				       int ifinterpolation )   // determine if do interpolation
{
	int m, i;
	double * approximation1, * detail1, * data1;

	if ( ifinterpolation == INTERPOLATION )
	{
		approximation1 = new double[n*2];
		detail1 = new double[n*2];
		data1 = new double[n*2];
		m = interpolation( approximation, n, approximation1 );
		m = interpolation( detail, n, detail1 );
	}
	else
	{
		approximation1 = new double[n];
		detail1 = new double[n];
		data1 = new double[n];
		for ( i = 0; i < n; i++ )
		{
			approximation1[i] = approximation[i];
			detail1[i] = detail[i];
		}
		m = n;
	}

	m = convolution( detail1, m, Hi, n_Hi, type, data, RECONSTRUCTION );
	m = convolution( approximation1, m, Lo, n_Lo, type, data1, RECONSTRUCTION );

	if ( ifinterpolation == INTERPOLATION )
	{
		for ( i = 0; i < m; i++ )
			data[i] = data1[i] + data[i];
	}
	else
	{
		for ( i = 0; i < m; i++ )
			data[i] = (data1[i] + data[i])/2;
	}

	delete approximation1;
	delete detail1;
	delete data1;

	return( m );
}
//////////////////////////////////////////////////////////////////////////
/*
 multi-level decomposition for 1D signal
*/
int multilevel_decomposition_1D
	( double * original,		// original signal
	int n,						// number of pointers, should be multiple of 2
	double * resultdata,		// decomposition data
	int nLevels,				// decomposition nLevels
	double * Hi, int n_Hi,		// high pass filter
	double * Lo, int n_Lo		// low pass filter
	)
{
	int i, x, nn;
	int rtv;
	double * l, * h, * data;

	nn = (n>>nLevels)<<nLevels; // guarantee nn to be multiple of 2
	for ( x = nn; x < n; x++ ) resultdata[x] = 0;
	l = new double[n];
	h = new double[n];
	data = new double[n];
	for ( x = 0; x < nn; x++ ) l[x] = original[x];
	for ( i = 0; i < nLevels; i++ )
	{
		for ( x = 0; x < nn; x++ ) data[x] = l[x];
		rtv = wavelet_analysis( data, nn, Hi, n_Hi, Lo, n_Lo, l, h, PERIODIC_EXTENSION, SUBSAMPLING);
		nn = nn/2;
		for ( x = 0; x < nn; x++ ) resultdata[x] = l[x];
		for ( x = 0; x < nn; x++ ) resultdata[nn+x] = h[x];
	}
	delete l;
	delete h;
	delete data;

	return( nn ); // return total points subject to decomposition
}
//////////////////////////////////////////////////////////////////////////
/*
 multi-level reconstruction for 1D signal
*/
int multilevel_reconstruction_1D
	( double * de_data,			// reconstructed signal
	int n,						// number of pointers, should be multiple of 2
	double * resultdata,		// reconstructed data
	int nLevels,					// reconstruct nLevels 
	double * Hi, int n_Hi,		// high pass filter
	double * Lo, int n_Lo		// low pass filter
	)
{
	int i, x, nn;
	int rtv;
	double * l, * h, * data;

	nn = n>>nLevels;
	l = new double[n];
	h = new double[n];
	data = new double[n];
	for ( x = 0; x < n; x++ ) data[x] = de_data[x];
	for ( i = 0; i < nLevels; i++ )
	{
		for ( x = 0; x < nn; x++ ) l[x] = data[x];
		for ( x = 0; x < nn; x++ ) h[x] = data[nn+x];
		rtv = wavelet_synthesis( l, h, nn, Hi, n_Hi, Lo, n_Lo, data, PERIODIC_EXTENSION, INTERPOLATION );
		nn = nn*2;
	}
	for ( x = 0; x < n; x++ ) resultdata[x] = data[x];
	delete l;
	delete h;
	delete data;

	return( nn ); // return total points subject to decomposition
}
//////////////////////////////////////////////////////////////////////////
/*
 multi-level decomposition for 2D signal
*/
int multilevel_decomposition_2D
	( double * image,			// original signal
	int width, int height,		// number of pointers, should be multiple of 2
	double * wavelet_data,			// decompositions data
	int nLevels,					// decomposition nLevels  
	double * Hi, int n_Hi,		// high pass filter
	double * Lo, int n_Lo		// low pass filter
	)
{
	int i, x, y, ww, hh;
	int rtv;
	double * l, * h, * data;

	ww = ( width>>nLevels )<<nLevels; // guarantee nn to be multiple of 2
	hh = ( height>>nLevels )<<nLevels;

	for ( y = 0; y < height; y++ )
		for ( x = 0; x < width; x++ )
			wavelet_data[y*width+x] = image[y*width+x];

	for ( i = 0; i < nLevels; i++ )
	{
		for ( y = 0; y < hh; y++ )
		{
			l = new double[ww];
			h = new double[ww];
			data = new double[ww];
			for ( x = 0; x < ww; x++ ) data[x] = wavelet_data[ y*width+x ];
			rtv = wavelet_analysis( data, ww, Hi, n_Hi, Lo, n_Lo, l, h, PERIODIC_EXTENSION, SUBSAMPLING);
			for ( x = 0; x < ww/2; x++ ) wavelet_data[y*width+x] = l[x];
			for ( x = 0; x < ww/2; x++ ) wavelet_data[y*width+x+ww/2] = h[x];
			delete l;
			delete h;
			delete data;
		}
		for ( x = 0; x < ww; x++ )
		{
			l = new double[hh];
			h = new double[hh];
			data = new double[hh];
			for ( y = 0; y < hh; y++ ) data[y] = wavelet_data[ y*width+x ];
			rtv = wavelet_analysis( data, hh, Hi, n_Hi, Lo, n_Lo, l, h, PERIODIC_EXTENSION, SUBSAMPLING);
			for ( y = 0; y < hh/2; y++ ) wavelet_data[y*width+x] = l[y];
			for ( y = 0; y < hh/2; y++ ) wavelet_data[(y+hh/2)*width+x] = h[y];
			delete l;
			delete h;
			delete data;
		}
		ww = ww/2;
		hh = hh/2;
	}

	return( nLevels ); // return decomposition nLevels
}
//////////////////////////////////////////////////////////////////////////
/*
 multi-level reconstruction for 2D signal
*/
int multilevel_reconstruction_2D
	( double * wavelet_data,	// wavelet coefficients
	int width, int height,		// number of pointers, should be multiple of 2
	double * image,				// reconstructed data
	int nLevels,					// reconstruction nLevels
	double * Hi, int n_Hi,		// high pass filter
	double * Lo, int n_Lo		// low pass filter
	)
{
	int i, x, y, ww, hh;
	int rtv;
	double * l, * h, * data;

	ww = width>>nLevels;  // guarantee nn to be multiple of 2
	hh = height>>nLevels;

	for ( y = 0; y < height; y++ )
		for ( x = 0; x < width; x++ )
			image[y*width+x] = wavelet_data[y*width+x];

	for ( i = 0; i < nLevels; i++ )
	{

		for ( x = 0; x < ww*2; x++ )
		{
			l = new double[hh];
			h = new double[hh];
			data = new double[hh*2];
			for ( y = 0; y < hh; y++ ) 
			{
				l[y] = image[ y*width+x ];
				h[y] = image[ (y+hh)*width+x ];
			}
			rtv = wavelet_synthesis( l, h, hh, Hi, n_Hi, Lo, n_Lo, data, PERIODIC_EXTENSION, INTERPOLATION );
			for ( y = 0; y < hh*2; y++ ) image[y*width+x] = data[y];
			delete l;
			delete h;
			delete data;
		}
		for ( y = 0; y < hh*2; y++ )
		{
			l = new double[ww];
			h = new double[ww];
			data = new double[ww*2];
			for ( x = 0; x < ww; x++ )
			{
				l[x] = image[ y*width+x ];
				h[x] = image[ y*width+x+ww ];
			}
			rtv = wavelet_synthesis( l, h, ww, Hi, n_Hi, Lo, n_Lo, data, PERIODIC_EXTENSION, INTERPOLATION );
			for ( x = 0; x < ww*2; x++ ) image[y*width+x] = data[x];
			delete l;
			delete h;
			delete data;
		}
		ww = ww*2;
		hh = hh*2;
	}

	return( nLevels ); // return decomposition nLevels
}
//////////////////////////////////////////////////////////////////////////
/*
 redundant one-level decomposition for 2D signal
*/
int decomposition_2D_r( double * image,         // original signal
					    int width, int height,  // number of pointers, should be multiple of 2
					    double * Hi, int n_Hi,  // high pass filter
					    double * Lo, int n_Lo,  // low pass filter
					    double * approximation, // decomposition data
					    double * detail_h,      // horizontal
						double * detail_v,      // vertical 
						double * detail_d )     // diagonal 
{
	int x, y, ww, hh;
	int rtv;
	double * l, * h, * data;

	ww = ( width>>1 )<<1; // guarantee nn to be multiple of 2
	hh = ( height>>1 )<<1;

	for ( y = 0; y < height; y++ )
		for ( x = 0; x < width; x++ )
			approximation[y*width+x] = image[y*width+x];

	for ( y = 0; y < hh; y++ )
	{
		l = new double[ww];
		h = new double[ww];
		data = new double[ww];
		for ( x = 0; x < ww; x++ ) 
			data[x] = approximation[ y*width+x ];
		rtv = wavelet_analysis( data, ww, Hi, n_Hi, Lo, n_Lo, l, h, PERIODIC_EXTENSION, NO_SUBSAMPLING);
		for ( x = 0; x < ww; x++ ) approximation[y*width+x] = l[x];
		for ( x = 0; x < ww; x++ ) detail_h[y*width+x] = h[x];
		delete l;
		delete h;
		delete data;
	}
	for ( x = 0; x < ww; x++ )
	{
		l = new double[hh];
		h = new double[hh];
		data = new double[hh];
		for ( y = 0; y < hh; y++ ) data[y] = approximation[ y*width+x ];
		rtv = wavelet_analysis( data, hh, Hi, n_Hi, Lo, n_Lo, l, h, PERIODIC_EXTENSION, NO_SUBSAMPLING);
		for ( y = 0; y < hh; y++ ) approximation[y*width+x] = l[y];
		for ( y = 0; y < hh; y++ ) detail_v[y*width+x] = h[y];
		delete l;
		delete h;
		delete data;
	}
	for ( x = 0; x < ww; x++ )
	{
		l = new double[hh];
		h = new double[hh];
		data = new double[hh];
		for ( y = 0; y < hh; y++ ) data[y] = detail_h[ y*width+x ];
		rtv = wavelet_analysis( data, hh, Hi, n_Hi, Lo, n_Lo, l, h, PERIODIC_EXTENSION, NO_SUBSAMPLING);
		for ( y = 0; y < hh; y++ ) detail_h[y*width+x] = l[y];
		for ( y = 0; y < hh; y++ ) detail_d[y*width+x] = h[y];
		delete l;
		delete h;
		delete data;
	}

	return( 1 ); // return decomposition nLevels
}
//////////////////////////////////////////////////////////////////////////
/*
 redundant one-level reconstruction for 2D signal
*/
int reconstruction_2D_r( double * approximation, // wavelet coefficients
					     double * detail_h,
						 double * detail_v,
						 double * detail_d,
						 int width, int height,  // number of pointers, should be multiple of 2
						 double * Hi, int n_Hi,  // high pass filter
						 double * Lo, int n_Lo,  // low pass filter
						 double * image )        // reconstructed data						
{
	int x, y, ww, hh;
	int rtv;
	double * l, * h, * data;

	ww = ( width>>1 )<<1; // guarantee nn to be multiple of 2
	hh = ( height>>1 )<<1;

	for ( y = 0; y < height; y++ )
		for ( x = 0; x < width; x++ )
			image[y*width+x] = approximation[y*width+x];

	for ( x = 0; x < ww; x++ )
	{
		l = new double[hh];
		h = new double[hh];
		data = new double[hh];
		for ( y = 0; y < hh; y++ ) 
		{
			l[y] = image[ y*width+x ];
			h[y] = detail_v[ y*width+x ];
		}
		rtv = wavelet_synthesis( l, h, hh, Hi, n_Hi, Lo, n_Lo, data, PERIODIC_EXTENSION, NO_INTERPOLATION );
		for ( y = 0; y < hh; y++ ) image[y*width+x] = data[y];
		delete l;
		delete h;
		delete data;
	}
	for ( x = 0; x < ww; x++ )
	{
		l = new double[hh];
		h = new double[hh];
		data = new double[hh];
		for ( y = 0; y < hh; y++ ) 
		{
			l[y] = detail_h[ y*width+x ];
			h[y] = detail_d[ y*width+x ];
		}
		rtv = wavelet_synthesis( l, h, hh, Hi, n_Hi, Lo, n_Lo, data, PERIODIC_EXTENSION, NO_INTERPOLATION );
		for ( y = 0; y < hh; y++ ) detail_h[y*width+x] = data[y];
		delete l;
		delete h;
		delete data;
	}
	for ( y = 0; y < hh; y++ )
	{
		l = new double[ww];
		h = new double[ww];
		data = new double[ww];
		for ( x = 0; x < ww; x++ )
		{
			l[x] = image[ y*width+x ];
			h[x] = detail_h[ y*width+x ];
		}
		rtv = wavelet_synthesis( l, h, ww, Hi, n_Hi, Lo, n_Lo, data, PERIODIC_EXTENSION, NO_INTERPOLATION );
		for ( x = 0; x < ww; x++ ) image[y*width+x] = data[x];
		delete l;
		delete h;
		delete data;
	}

	return(1); // return decomposition nLevels
}
//////////////////////////////////////////////////////////////////////////
/*
 
*/
Mat getForwardWaveletBasis(Mat img, int nLevels)
{
	int nRows = img.rows;
	int nCols = img.cols;
	Mat srcImg(nRows, nCols, CV_64FC1);
	Mat dstImg(nRows, nCols, CV_64FC1);
	Mat wavImg(nRows, nCols, CV_64FC1);
	double *src = new double[nRows * nCols];
	double *wav = new double[nRows * nCols];
	double *dst = new double[nRows * nCols];
	for(int i=0; i<nRows; i++)
	{
		for(int j=0; j<nCols; j++)
		{
			srcImg.at<double>(i, j) = img.at<double>(i, j);
			src[i*nCols + j] = srcImg.at<double>(i, j);
		}
	}
	multilevel_decomposition_2D(src, nCols, nRows, dst, nLevels, 
								Hd, 4, Ld, 4);
	for(int i=0; i<nRows; i++)
	{
		for(int j=0; j<nCols; j++)
		{
			dstImg.at<double>(i, j) = dst[i*nCols + j];
		}
	}
	divide(dstImg, srcImg, wavImg);
	srcImg.release();
	dstImg.release();
	delete src;
	delete wav;
	delete dst;
	return wavImg;
}
//////////////////////////////////////////////////////////////////////////
/*
 
*/
Mat getInverseWaveletBasis(Mat img, int nLevels)
{
	int nRows = img.rows;
	int nCols = img.cols;
	Mat srcImg(nRows, nCols, CV_64FC1);
	Mat dstImg(nRows, nCols, CV_64FC1);
	Mat wavImg(nRows, nCols, CV_64FC1);
	double *src = new double[nRows * nCols];
	double *wav = new double[nRows * nCols];
	double *dst = new double[nRows * nCols];
	for(int i=0; i<nRows; i++)
	{
		for(int j=0; j<nCols; j++)
		{
			srcImg.at<double>(i, j) = img.at<double>(i, j);
			src[i*nCols + j] = srcImg.at<double>(i, j);
		}
	}
	multilevel_reconstruction_2D(src, nCols, nRows, dst, nLevels, 
								Hr, 4, Lr, 4);
	for(int i=0; i<nRows; i++)
	{
		for(int j=0; j<nCols; j++)
		{
			dstImg.at<double>(i, j) = dst[i*nCols + j];
		}
	}
	divide(dstImg, srcImg, wavImg);
	srcImg.release();
	dstImg.release();
	delete src;
	delete wav;
	delete dst;
	return wavImg;
}
//////////////////////////////////////////////////////////////////////////
/*
 
*/
Mat getForwardWaveletBasis2(Mat img, int nLevels)
{
	Mat planes_img[2];
	Mat planes_fdwt[2];

	split(img, planes_img);
	planes_fdwt[0] = getForwardWaveletBasis(planes_img[0], nLevels);
	planes_fdwt[1] = getForwardWaveletBasis(planes_img[1], nLevels);

	Mat fdwt;
	merge(planes_fdwt, 2, fdwt);

	return fdwt;
}
//////////////////////////////////////////////////////////////////////////
/*
 
*/
Mat getInverseWaveletBasis2(Mat img, int nLevels)
{
	Mat planes_img[2];
	Mat planes_idwt[2];

	split(img, planes_img);
	planes_idwt[0] = getInverseWaveletBasis(planes_img[0], nLevels);
	planes_idwt[1] = getInverseWaveletBasis(planes_img[1], nLevels);

	Mat idwt;
	merge(planes_idwt, 2, idwt);

	return idwt;
}